package com.chb.dmp.etl

import org.apache.spark.sql.types.StringType
import org.apache.spark.sql.{Dataset, Row}
import org.lionsoul.ip2region.{DbConfig, DbSearcher}

object IPProcessor extends Processor {
  override def process(ds: Dataset[Row]): Dataset[Row] = {
    // 原始数据集 c1, c2 ... ip
    // 新数据集   c1, c2 ... ip city

    // 如何转换
    // 方法一: 生成只有IP的数据集(IP, city), 和原来数据集按照IP 进行 join
    // 方法二： 直接插入四列？ withColumn, 编写UDF
    // 方法三： 生成新的数据，然后修改原来的schema

    // 此处使用第三种方法，因为修改schema, 也是一个比较常见的操作

    // 1、 将ds转为一个新的数据集 添加数据city
   val convertRdd =  ds.rdd.mapPartitions(convertIP)
    // 2、生成新的schema
    val newSchema = ds.schema.add("city", StringType)
    // 3、使用RDD  加上新的schema, 生成新的Dataset
    val newDs: Dataset[Row] = ds.sparkSession.createDataFrame(convertRdd, newSchema)

    newDs
  }


  /**
   * 将ip转为城市信息
   *
   * @param iterator
   * @return
   */
  def convertIP(iterator: Iterator[Row]): Iterator[Row] = {
    val searcher = new DbSearcher(new DbConfig(), "dataset/ip2region.db")

    iterator.map(row => {
      val ip = row.getAs[String]("ip")
      val region = searcher.btreeSearch(ip).getRegion
      val city = region.split("\\|")(3)

      val rowSeq = row.toSeq :+ city
      Row.fromSeq(rowSeq)
    })
  }


}
